#ifndef HIXL_CS_CLIENT_H
#define HIXL_CS_CLIENT_H

#include "hccl_api.h"
#include "EndpointStore.h"
#include "hccl_socket.h"
#include "HixlMemStore.h"
#include <map>
#include <vector>
#include <mutex>
#include <memory>
#include <atomic>
#include <string>


namespace hixl {
    using ClientHandle = void *;
    using EndpointHandle = void *;
    using ChannelHandle = void *;

    enum class ConnectPhase {
        IDLE,
        RA_CONNECTING,
        HANDSHAKE_SEND_EP,
        HANDSHAKE_RECV_EP,
        CHANNEL_CREATING,
        DONE,
        FAILED
    };

    struct EndpointMemInfo {
        hccl::HcclMem mem_info;
        char *mem_tag;
        EndpointHandle endpoint_handle;
    };

    struct EndpointChannelInfo {
        EndpointHandle endpoint_handle{nullptr};
        ChannelHandle channel_handle{nullptr};
        std::shared_ptr<hccl::HcclSocket> socket;
        bool is_connected{false};
    };

    /**
     * @brief 单边通信 Client 实现类
     */
    class HixlCsClient {
    public:
        HixlCsClient();

        ~HixlCsClient();

        hccl::HcclResult Create(char *server_ip,
                                uint32_t server_port,
                                const hccl::EndPoint *end_point);

        hccl::HcclResult Connect();

        hccl::HcclResult GetStatus(int32_t *status);

        hccl::HcclResult GetRemoteMem(hccl::HcclMem **remote_mem_list,
                                      char **mem_tag_list,
                                      uint32_t *list_num,
                                      uint32_t timeout);
        
        // 注册client的endpoint的内存信息到内存注册表中。                              
        hccl::HcclResult ClientRegMem(const char *mem_tag, const HcclMem *mem, MemHandle *mem_handle);

        //通过已经建立好的channel，从用户提取的地址列表中，批量读取server内存地址中的内容
        hccl::HcclResult BatchGet(uint32_t list_num, void** remote_buf_list, 
                            const void** local_buf_list, uint64_t* len_list);         
              
        //通过已经建立好的channel，从用户提取的地址列表中，批量往server内存地址中写入内容
        hccl::HcclResult BatchPut(uint32_t list_num, void** remote_buf_list, 
                            const void** local_buf_list, uint64_t* len_list);   
                            
        //通过已经建立好的channel，检查批量读写的状态。
        hccl::HcclResult CheckStatus();                     

        // 注册client的endpoint的内存信息到内存注册表中。                              
        hccl::HcclResult ClientUnRegMem();

        hccl::HcclResult Destroy();

        // 供事件循环调用，用于推进握手/建通道
        hccl::HcclResult OnSocketEvent(hccl::FdHandle fd);

    private:
        // === 创建/销毁相关 ===
        hccl::HcclResult InitState(const char *server_ip,
                                   uint32_t server_port,
                                   const hccl::EndPoint *end_point);

        hccl::HcclResult CreateSelfEndpoint(const hccl::EndPoint *end_point);

        hccl::HcclResult DestroyAllEndpointsAndChannels();

        void ClearRemoteMemInfo();

        void ResetConnectState();

        // === RA Socket & 状态机 ===
        hccl::HcclResult InitSocketIfNeeded();

        hccl::HcclResult StartRaConnect();

        hccl::HcclResult HandleRaConnecting(int32_t *status);

        hccl::HcclResult HandleSocketOk();

        hccl::HcclResult RegisterSocketToEventLoop();

        void UpdateStatusByPhase(int32_t *status) const;

        // === 握手 + Channel 创建 ===
        hccl::HcclResult AdvanceHandshakeAndChannel();

        hccl::HcclResult HandleHandshakeSendEp();

        hccl::HcclResult HandleHandshakeRecvEp();

        hccl::HcclResult HandleChannelCreating();

        EndpointChannelInfo *FindActiveChannelUnlocked();

        // === 远端内存获取 ===
        struct RemoteMemReqHeader {
            uint32_t magic;
            uint32_t version;
            uint32_t timeout;
            uint32_t reserved;
        };

        struct RemoteMemRequest {
            RemoteMemReqHeader header;
            hccl::EndPoint client_ep;
        };

        hccl::HcclResult GetActiveSocketAndEndpoint(
            std::shared_ptr<hccl::HcclSocket> &socket,
            EndpointHandle &endpoint);

        hccl::HcclResult SendRemoteMemRequest(hccl::HcclSocket &socket,
                                              uint32_t timeout);

        hccl::HcclResult RecvRemoteMemCount(hccl::HcclSocket &socket,
                                            uint32_t *memCount,
                                            hccl::HcclMem **remote_mem_list,
                                            char **mem_tag_list);

        hccl::HcclResult AllocateRemoteMemArrays(uint32_t memCount,
                                                 hccl::HcclMem **memList,
                                                 char **tagBufAll);

        hccl::HcclResult RecvOneMemEntry(hccl::HcclSocket &socket,
                                         EndpointHandle endpoint,
                                         hccl::HcclMem &outMem,
                                         char *tagSlot,
                                         uint32_t tagSlotLen);

        void CleanupRemoteMemOnError(hccl::HcclMem *memList,
                                     char *tagBufAll) const;

    private:
        static constexpr uint32_t TAG_MAX_LEN = 64;
        static constexpr uint32_t REMOTE_MAGIC = 0x524D4151u; // 'RMAQ'
        static constexpr uint32_t REMOTE_VER = 1;

        std::string server_ip_;
        uint32_t server_port_{0};

        hccl::EndPoint self_endpoint_{};
        EndpointHandle self_endpoint_handle_{nullptr};

        std::unique_ptr<EndpointStore> endpoint_store_;
        bool is_connected_{false};

        mutable std::mutex mutex_;
        std::map<uint64_t,
            std::vector<EndpointMemInfo> > remote_mem_info_;

        std::mutex chn_mutex_;
        std::map<hccl::FdHandle,
            EndpointChannelInfo> channels_;

        // 连接状态机
        std::shared_ptr<hccl::HcclSocket> socket_;
        ConnectPhase phase_{ConnectPhase::IDLE};
        bool connect_started_{false};

        hccl::EndPoint peer_endpoint_{};
        size_t send_ep_offset_{0};
        size_t recv_ep_offset_{0};
        hccl::FdHandle registered_fd_{nullptr};
    };
} // namespace hixl

#endif // HIXL_CS_CLIENT_H
